#!/usr/bin/env python
#
# Author: Qiming Sun <osirpt.sun@gmail.com>
#

'''
Active space can be adjusted by specifing the number of orbitals for each irrep.
'''

import tempfile
import h5py
import time
import numpy as np
import csv
import mo_sym
#import p2ftest

from pyscf import gto, scf, mcscf, ao2mo, cc, symm
def view(h5file, dataname='eri_mo'):
    f5 = h5py.File(h5file, 'r')
    print('dataset %s, shape %s' % (str(f5.keys()), str(f5[dataname].shape)))
    f5.close()

def symminfo(groupname):
   if groupname == 'Dooh':
      group_corr = 'D2h'
   elif groupname == 'Coov':
      group_corr = 'C2v'
   else:
      group_corr = mol.groupname
   symm_table = symm.param.IRREP_ID_TABLE[group_corr]
#   print('symm_table', symm_table)
   nirrep = len(symm_table)
   irrep_sym = ['']*nirrep
   for irsymb in symm_table:
      irid = symm_table[irsymb]
      irrep_sym[irid] = irsymb
#      print('irid', irid)
#   print('irrep_sym', irrep_sym)
   return group_corr, nirrep, irrep_sym

def symorb_sort(mol, nirrep, mo_occ, mo_coeff):
   listocc = list(np.argwhere(mo_occ > 0).flatten())
   listvir = list(np.argwhere(mo_occ == 0).flatten())
   nocc = len(listocc)
   nvir = len(listvir)
#   print('nocc, nvir', nocc, nvir)
   symname = symm.label_orb_symm(mol, mol.irrep_name, mol.symm_orb, mo_coeff)
   symid = symm.label_orb_symm(mol, mol.irrep_id, mol.symm_orb, mo_coeff)
   symid = symid % 10
#   print('irrep name:', symname)
#   print('irrep   id:', symid)
   nmo = nocc + nvir
   isympop = np.zeros((nmo,8), dtype = int)
   isymvrt = np.zeros((nmo,8), dtype = int)
   isymmo  = np.zeros((nmo,8), dtype = int)
   pop = np.zeros(8, dtype = int)
   vrt = np.zeros(8, dtype = int)
   for irrep in range(nirrep):
       irrep_occ = list(np.argwhere(symid[:nocc] == irrep).flatten())
       irrep_vir = list(np.argwhere(symid[nocc:] == irrep).flatten() + nocc)
       pop[irrep] = len(irrep_occ)
       vrt[irrep] = len(irrep_vir)
       isympop[0:pop[irrep],irrep]=irrep_occ
       isymmo[0:pop[irrep],irrep] = irrep_occ
       isymvrt[0:vrt[irrep],irrep]=irrep_vir
       isymmo[pop[irrep]:pop[irrep]+vrt[irrep],irrep] = irrep_vir
#      print('irrep, irrep_occ', irrep, irrep_occ, isympop[:,irrep])
#      print('irrep, irrep_vir', irrep, irrep_vir, isymvrt[:,irrep])

#   print('isymmo', isymmo)
   return pop, vrt, isympop, isymvrt, isymmo

mol = gto.Mole()
mol.build(
       atom = '''
  O                
  H    1    1.00000000   
  H    1    1.00000000   2 110.000000 
''',
    basis = 'cc-pvdz',
    symmetry = True,
    verbose = 5,
    charge = 0,
    spin = 0,
    output = 'outmo_py',
    max_memory = 40000
)
myhf = scf.UHF(mol)
myhf.kernel()
print('scf energy', myhf.e_tot)
#print(myhf.mo_energy)

time1 = time.time()
nocpa,nocpb = myhf.mo_occ
noa = sum(nocpa>0)
nob = sum(nocpb>0)
nva = sum(nocpa==0)
nvb = sum(nocpb==0)
nmo = noa+nva
escf = myhf.e_tot

group_corr, nirrep, irrep_sym = symminfo(mol.groupname)

print('noa, nob', noa, nob, nva, nvb, nmo)
with open('cc_pyscf_int', 'w') as csvfile:
    file_cc = csv.writer(csvfile)
    file_cc.writerow([noa, nob, nva, nvb, nmo, escf, group_corr])
    orb_energy = myhf.mo_energy
    file_cc.writerow(orb_energy[0])
    file_cc.writerow(orb_energy[1]) 

orb = myhf.mo_coeff
orba = orb[0]
orbb = orb[1]
print(orba.shape, orbb.shape)

time1 = time.time()
nspin = 2
pop = np.zeros((8,nspin), dtype = int)
vrt = np.zeros((8,nspin), dtype = int)

pop[:,0], vrt[:,0], isympopa, isymvrta, isymmo = symorb_sort(mol, nirrep, myhf.mo_occ[0], myhf.mo_coeff[0])

pop[:,1], vrt[:,1], isympopb, isymvrtb, isymmo = symorb_sort(mol, nirrep, myhf.mo_occ[1], myhf.mo_coeff[1])

print('groupname', group_corr)

dirprd = np.array(((0, 1, 2, 3, 4, 5, 6, 7),
          (1, 0, 3, 2, 5, 4, 7, 6),
          (2, 3, 0, 1, 6, 7, 4, 5),
          (3, 2, 1, 0, 7, 6, 5, 4),
          (4, 5, 6, 7, 0, 1, 2, 3),
          (5, 4, 7, 6, 1, 0, 3, 2),
          (6, 7, 4, 5, 2, 3, 0, 1),
          (7, 6, 5, 4, 3, 2, 1, 0)))
mo = (orba, orba, orba, orba)
vrta = vrt[:,0]
popa = pop[:,0]
vrtb = vrt[:,1]
popb = pop[:,1]
ioffvvvv, nvvvv, iposvv, numvv, iposvv, numvv = mo_sym.getpqrs_incore(dirprd, nirrep, 
                                                       vrta, vrta, 1, vrta, vrta, 1)
ioffvvvo, nvvvo, iposvv, numvv, iposvo, numvo = mo_sym.getpqrs_incore(dirprd, nirrep, 
                                                       vrta, vrta, 1, vrta, popa, 1)
ioffvvoo, nvvoo, iposvv, numvv, iposoo, numoo = mo_sym.getpqrs_incore(dirprd, nirrep, 
                                                       vrta, vrta, 1, popa, popa, 1)
ioffvovo, nvovo, iposvo, numvo, iposvo, numvo = mo_sym.getpqrs_incore(dirprd, nirrep, 
                                                       vrta, popa, 1, vrta, popa, 1)
ioffooov, nooov, iposoo, numoo, iposov, numov = mo_sym.getpqrs_incore(dirprd, nirrep, 
                                                       popa, popa, 1, popa, vrta, 1)
ioffoooo, noooo, iposoo, numoo, iposoo, numoo = mo_sym.getpqrs_incore(dirprd, nirrep, 
                                                        popa, popa, 1, popa, popa, 1)

mov = orba[:,noa:nmo]
moo = orba[:,0:noa]
movvvv = (mov, mov, mov, mov)
movvvo = (mov, mov, moo, mov)
movvoo = (mov, mov, moo, moo) 
movovo = (moo, mov, moo, mov) 
moooov = (moo, moo, mov, moo) 
mooooo = (moo, moo, moo, moo) 

mo_sym.save_symm(dirprd, nirrep, nmo, vrta, popa, isymmo)

ivrtin = isymvrta - noa
ipopin = isympopa

time1 = time.time()

erimo = ao2mo.general(mol, movvvv)
ivvvv = mo_sym.mo_pqrs2(erimo, dirprd, nirrep, nmo, numvv, numvv, iposvv, iposvv, ioffvvvv, 
                        vrta, vrta, vrta, vrta, nva, nva, nva, nva, 
                        ivrtin, ivrtin, ivrtin, ivrtin,
                        0, 0, nva*(nva+1)/2, nva*(nva+1)/2, nvvvv, 'VVVV_SYM')
print('norm of vvvv', nvvvv, np.linalg.norm(ivvvv))

erimo = ao2mo.general(mol, movvvo)
ivvvo = mo_sym.mo_pqrs2(erimo, dirprd, nirrep, nmo, numvv, numvo, iposvv, iposvo, ioffvvvo, 
                        vrta, vrta, vrta, popa, nva, nva, nva, noa, 
                        ivrtin, ivrtin, ivrtin, ipopin,
                        0, 1, nva*(nva+1)/2, nva*noa, nvvvo, 'VVVO_SYM')
print('norm of vvvo', nvvvo, np.linalg.norm(ivvvo))

erimo = ao2mo.general(mol, movvoo)
ivovo = mo_sym.mo_pqrs2(erimo, dirprd, nirrep, nmo, numvo, numvo, iposvo, iposvo, ioffvovo, 
                        vrta, popa, vrta, popa, nva, noa, nva, noa, 
                        ivrtin, ipopin, ivrtin, ipopin,
                        0, 0, nva*(nva+1)/2, noa*(noa+1)/2, nvovo, 'VOVO_SYM')
print('norm of vovo', nvvoo, np.linalg.norm(ivovo))

erimo = ao2mo.general(mol, movovo)
ivvoo = mo_sym.mo_pqrs2(erimo, dirprd, nirrep, nmo, numvv, numoo, iposvv, iposoo, ioffvvoo, 
                        vrta, vrta, popa, popa, nva, nva, noa, noa, 
                        ivrtin, ivrtin, ipopin, ipopin,
                        1, 1, nva*noa, nva*noa, nvvoo, 'VVOO_SYM')
print('norm of vvoo', nvvoo, np.linalg.norm(ivvoo))

erimo = ao2mo.general(mol, moooov)
iooov = mo_sym.mo_pqrs2(erimo, dirprd, nirrep, nmo, numoo, numvo, iposoo, iposov, ioffooov, 
                        popa, popa, popa, vrta, noa, noa, noa, nva,
                        ipopin, ipopin, ipopin, ivrtin,
                        0, 1, noa*(noa+1)/2, nva*noa, nooov, 'OOOV_SYM')
print('norm of ooov', nooov, np.linalg.norm(iooov))

erimo = ao2mo.general(mol, mooooo)
ioooo = mo_sym.mo_pqrs2(erimo, dirprd, nirrep, nmo, numoo, numoo, iposoo, iposoo, ioffoooo, 
                        popa, popa, popa, popa, noa, noa, noa, noa,
                        ipopin, ipopin, ipopin, ipopin,
                        0, 0, noa*(noa+1)/2, noa*(noa+1)/2, noooo, 'OOOO_SYM')
print('norm of ooov', noooo, np.linalg.norm(ioooo))

time2 = time.time()
print(' time for integral transformation: ', time2-time1)

